package com.cyzy.chatai.socket;

import com.cyzy.chatai.service.AiChatService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket服务器处理类
 */
@Component
public class ChatWebSocket implements WebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(ChatWebSocket.class);
    private final SimpleDateFormat format = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
    
    private final AiChatService aiChatService;
    
    private final ConcurrentHashMap<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    // 构造器注入
    public ChatWebSocket(AiChatService aiChatService) {
        this.aiChatService = aiChatService;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.put(session.getId(), session);
        logger.info("WebSocket连接建立成功，会话ID: {}，当前连接数: {}", session.getId(), sessions.size());
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        logger.info("接收到WebSocket消息 - 会话ID: {}，消息: {}", session.getId(), message.getPayload());
        String msg = message.getPayload().toString();

        try {
            // 调用AI服务获取回复
            String aiReply = aiChatService.sendMessage(msg, session.getId());
            if (session.isOpen()) {
                session.sendMessage(new TextMessage(aiReply));
            }
        } catch (Exception e) {
            logger.error("处理WebSocket消息异常 - 会话ID: {}: {}", session.getId(), e.getMessage(), e);
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage("抱歉，AI服务暂时不可用，请稍后重试。"));
                }
            } catch (IOException ioException) {
                logger.error("发送错误消息失败 - 会话ID: {}: {}", session.getId(), ioException.getMessage());
            }
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        logger.error("WebSocket传输异常 - 会话ID: {}: {}", session.getId(), exception.getMessage(), exception);
        sessions.remove(session.getId());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        sessions.remove(session.getId());
        logger.info("WebSocket连接关闭 - 会话ID: {}，状态: {}，当前连接数: {}", session.getId(), closeStatus.getCode(), sessions.size());
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    public ConcurrentHashMap<String, WebSocketSession> getSessions() {
        return sessions;
    }
    
    public int getConnectionCount() {
        return sessions.size();
    }
}
